#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
图片搜索组件
提供多种图片搜索API的集成，支持Google Images、Bing Images等
"""

import os
import json
import time
import hashlib
from datetime import datetime
from urllib.parse import urlencode, quote
from urllib import parse
from concurrent.futures import ThreadPoolExecutor
import threading
from PyQt5.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QPushButton,
                             QLabel, QLineEdit, QComboBox, QSpinBox, QGroupBox,
                             QGridLayout, QTextEdit, QCheckBox, QSlider)
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QTimer
from PyQt5.QtGui import QFont
import requests
from bs4 import BeautifulSoup
import re

class ImageSearchAPI:
    """图片搜索API基类"""
    
    def __init__(self, name):
        self.name = name
        self.session = requests.Session()
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        })
        
    def search(self, query, page=1, per_page=20, **kwargs):
        """搜索图片"""
        raise NotImplementedError
        
    def parse_image_data(self, raw_data):
        """解析图片数据"""
        raise NotImplementedError

class BingImagesAPI(ImageSearchAPI):
    """Bing Images API - 使用与百度图片搜索相同的处理逻辑"""
    
    def __init__(self):
        super().__init__("Bing Images")
        # 使用与百度图片搜索相同的URL模板和请求头
        self.url_template = 'https://www.bing.com/images/search?q={}&form=HDRSC2&first={}&count={}'
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.36',
            'Referer': 'https://www.bing.com',
            'Accept': '*/*',
            'Accept-Language': 'zh-CN,zh;q=0.9',
            'Connection': 'keep-alive'
        })
        
    def search(self, query, page=1, per_page=20, **kwargs):
        """搜索Bing Images - 使用与百度图片搜索相同的处理逻辑"""
        try:
            # 使用URL模板构建请求
            start = (page - 1) * per_page
            url = self.url_template.format(
                parse.quote(query),  # 搜索关键词
                start + 1,           # first参数
                per_page             # count参数
            )
            
            print(f"\n🔍 [Bing图片搜索] 请求URL: {url}")
            print(f"📝 [Bing图片搜索] 搜索关键词: {query}")
            print(f"📄 [Bing图片搜索] 页码: {page}, 每页数量: {per_page}")
            print(f"🌐 [Bing图片搜索] 请求头: {dict(self.session.headers)}")
            
            response = self.session.get(url, timeout=15)
            print(f"📊 [Bing图片搜索] 响应状态码: {response.status_code}")
            print(f"📏 [Bing图片搜索] 响应内容长度: {len(response.content)} bytes")
            
            response.raise_for_status()
            
            # 解析HTML获取图片数据
            soup = BeautifulSoup(response.text, 'html.parser')
            print(f"📋 [Bing图片搜索] HTML解析成功，找到 {len(soup.find_all('img', class_='mimg'))} 个mimg类图片")
            
            images = []
            
            # 查找图片元素 - 使用与百度图片搜索相同的逻辑
            img_elements = soup.find_all('img', class_='mimg')
            
            for i, img in enumerate(img_elements):
                # 提取图片信息 - 使用与百度图片搜索相同的字段名
                image_info = self._extract_image_info(img)
                if image_info:
                    images.append(image_info)
                    print(f"✅ [Bing图片搜索] 成功提取图片 {len(images)}: {image_info.get('title', '无标题')[:30]}...")
                else:
                    print(f"❌ [Bing图片搜索] 提取图片信息失败: {img}")
                    
            print(f"🎯 [Bing图片搜索] 总共找到 {len(images)} 张有效图片")
            return images
            
        except Exception as e:
            print(f"❌ [Bing图片搜索] 搜索错误: {e}")
            print(f"🔍 [Bing图片搜索] 错误类型: {type(e).__name__}")
            return []
            
    def _extract_image_info(self, img):
        """提取图片信息 - 使用与百度图片搜索相同的逻辑"""
        try:
            # 获取缩略图URL
            src = img.get('src') or img.get('data-src')
            if not src or not src.startswith('http'):
                return None
                
            # 获取图片标题
            title = img.get('alt', '') or img.get('title', '') or ''
            if not title:
                title = 'Bing图片'
                
            # 获取图片尺寸
            width = img.get('width', 300)
            height = img.get('height', 200)
            
            # 确保尺寸是数字
            try:
                width = int(width) if width else 300
                height = int(height) if height else 200
            except (ValueError, TypeError):
                width, height = 300, 200
                
            return {
                'url': src,
                'title': title,
                'thumbnail': src,
                'width': width,
                'height': height,
                'source': 'Bing Images'
            }
            
        except Exception as e:
            print(f"提取图片信息错误: {e}")
            return None

class BaiduImagesAPI(ImageSearchAPI):
    """百度图片搜索API"""
    
    def __init__(self):
        super().__init__("百度图片")
        # 使用更稳定的URL模板
        self.url_template = 'https://image.baidu.com/search/acjson?tn=resultjson_com&ipn=rj&queryWord={}&word={}&pn={}&rn={}'
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.36',
            'Referer': 'https://image.baidu.com',
            'Accept': '*/*',
            'Accept-Language': 'zh-CN,zh;q=0.9',
            'Connection': 'keep-alive'
        })
        
    def search(self, query, page=1, per_page=20, **kwargs):
        """搜索百度图片"""
        try:
            # 使用URL模板构建请求
            start = (page - 1) * per_page
            url = self.url_template.format(
                parse.quote(query),  # queryWord
                parse.quote(query),  # word
                start,               # pn
                per_page             # rn
            )
            
            print(f"\n🔍 [百度图片搜索] 请求URL: {url}")
            print(f"📝 [百度图片搜索] 搜索关键词: {query}")
            print(f"📄 [百度图片搜索] 页码: {page}, 每页数量: {per_page}")
            print(f"🌐 [百度图片搜索] 请求头: {dict(self.session.headers)}")
            
            response = self.session.get(url, timeout=15)
            print(f"📊 [百度图片搜索] 响应状态码: {response.status_code}")
            print(f"📏 [百度图片搜索] 响应内容长度: {len(response.content)} bytes")
            
            response.raise_for_status()
            
            # 解析JSON数据
            data = response.json()
            print(f"📋 [百度图片搜索] JSON解析成功，数据类型: {type(data)}")
            
            if 'data' in data:
                print(f"📊 [百度图片搜索] 数据项数量: {len(data['data']) if isinstance(data['data'], list) else '非列表'}")
            else:
                print("⚠️ [百度图片搜索] 响应中没有'data'字段")
                print(f"🔍 [百度图片搜索] 响应内容: {data}")
            
            images = []
            
            if 'data' in data and isinstance(data['data'], list):
                for i, item in enumerate(data['data']):
                    if not item or not isinstance(item, dict):
                        print(f"⚠️ [百度图片搜索] 跳过无效数据项 {i}: {item}")
                        continue
                        
                    # 提取图片信息
                    image_info = self._extract_image_info(item)
                    if image_info:
                        images.append(image_info)
                        print(f"✅ [百度图片搜索] 成功提取图片 {len(images)}: {image_info.get('title', '无标题')[:30]}...")
                    else:
                        print(f"❌ [百度图片搜索] 提取图片信息失败: {item}")
                        
            print(f"🎯 [百度图片搜索] 总共找到 {len(images)} 张有效图片")
            return images
            
        except Exception as e:
            print(f"❌ [百度图片搜索] 搜索错误: {e}")
            print(f"🔍 [百度图片搜索] 错误类型: {type(e).__name__}")
            return []
            
    def _extract_image_info(self, item):
        """提取图片信息"""
        try:
            # 获取缩略图URL
            thumb_url = item.get('thumbURL', '')
            if not thumb_url:
                return None
                
            # 获取原始图片URL（百度可能返回加密的URL，我们使用缩略图作为备用）
            obj_url = item.get('objURL', '')
            if not obj_url or obj_url.startswith('ipprf_'):
                obj_url = thumb_url
                
            # 获取图片标题
            title = item.get('fromPageTitle', '') or item.get('title', '') or ''
            if not title:
                title = '百度图片'
                
            # 获取图片尺寸
            width = item.get('width', 300)
            height = item.get('height', 200)
            
            # 确保尺寸是数字
            try:
                width = int(width) if width else 300
                height = int(height) if height else 200
            except (ValueError, TypeError):
                width, height = 300, 200
                
            return {
                'url': obj_url,
                'title': title,
                'thumbnail': thumb_url,
                'width': width,
                'height': height,
                'source': '百度图片'
            }
            
        except Exception as e:
            print(f"提取图片信息错误: {e}")
            return None

class ImageSearchWorker(QThread):
    """图片搜索工作线程"""
    search_finished = pyqtSignal(list, str)  # 搜索结果, 错误信息
    progress_updated = pyqtSignal(int)  # 进度更新
    
    def __init__(self, query, search_engine, max_pages=1, per_page=20, start_page=1, **kwargs):
        super().__init__()
        self.query = query
        self.search_engine = search_engine
        self.max_pages = max_pages  # 最大搜索页数
        self.per_page = per_page
        self.start_page = start_page  # 起始页码
        self.kwargs = kwargs
        
    def run(self):
        """执行多页搜索"""
        try:
            self.progress_updated.emit(10)
            
            # 创建搜索API实例
            api = self._create_api()
            if not api:
                self.search_finished.emit([], "不支持的搜索引擎")
                return
                
            self.progress_updated.emit(20)
            
            all_images = []
            total_pages = self.max_pages
            
            # 多页搜索
            for page in range(self.start_page, self.start_page + self.max_pages):
                # 检查线程是否被请求停止
                if self.isInterruptionRequested():
                    print("🛑 [ImageSearchWorker] 线程被请求停止")
                    return
                    
                try:
                    print(f"\n📄 [多页搜索] 正在搜索第 {page}/{self.start_page + self.max_pages - 1} 页...")
                    
                    # 执行搜索
                    images = api.search(
                        self.query, 
                        page=page, 
                        per_page=self.per_page,
                        **self.kwargs
                    )
                    
                    if images:
                        all_images.extend(images)
                        print(f"✅ [多页搜索] 第 {page} 页找到 {len(images)} 张图片")
                    else:
                        print(f"⚠️ [多页搜索] 第 {page} 页没有找到图片，停止搜索")
                        break
                        
                    # 更新进度
                    progress = 20 + ((page - self.start_page + 1) * 60 // self.max_pages)
                    self.progress_updated.emit(progress)
                    
                    # 添加延迟避免请求过于频繁
                    if page < self.start_page + self.max_pages - 1:
                        self.msleep(1000)  # 1秒延迟
                        
                except Exception as e:
                    print(f"❌ [多页搜索] 第 {page} 页搜索失败: {e}")
                    continue
            
            self.progress_updated.emit(80)
            
            # 处理结果
            processed_images = self._process_images(all_images)
            
            print(f"🎯 [多页搜索] 总共搜索了 {self.max_pages} 页，找到 {len(all_images)} 张原始图片，处理后 {len(processed_images)} 张有效图片")
            
            self.progress_updated.emit(100)
            self.search_finished.emit(processed_images, "")
            
        except Exception as e:
            self.search_finished.emit([], str(e))
            
    def _create_api(self):
        """创建API实例"""
        if self.search_engine == "Bing Images":
            return BingImagesAPI()
        elif self.search_engine == "百度图片":
            return BaiduImagesAPI()
        else:
            return None
            
    def _process_images(self, images):
        """处理图片数据"""
        processed = []
        print(f"🔍 [ImageSearchWorker] 开始处理 {len(images)} 张图片")
        
        for i, img in enumerate(images):
            # 验证图片URL
            if self._is_valid_image_url(img['url']):
                processed.append(img)
                print(f"✅ [ImageSearchWorker] 图片 {i+1} 通过验证: {img.get('title', '无标题')[:30]}...")
            else:
                print(f"❌ [ImageSearchWorker] 图片 {i+1} 未通过验证: {img['url'][:50]}...")
                
        print(f"🎯 [ImageSearchWorker] 处理完成，有效图片: {len(processed)}/{len(images)}")
        return processed
        
    def _is_valid_image_url(self, url):
        """验证图片URL"""
        if not url or not url.startswith('http'):
            print(f"❌ [URL验证] 无效URL: {url}")
            return False
            
        # 检查文件扩展名
        valid_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp']
        url_lower = url.lower()
        
        # 如果URL包含查询参数，检查路径部分
        if '?' in url:
            path = url.split('?')[0]
        else:
            path = url
            
        # 检查扩展名
        has_extension = any(path.endswith(ext) for ext in valid_extensions)
        
        # 检查是否包含图片相关关键词
        has_image_keyword = any(keyword in url_lower for keyword in ['image', 'img', 'photo', 'pic'])
        
        # 对于百度图片和Bing图片，放宽验证条件
        is_baidu_image = 'baidu.com' in url_lower and ('it/u=' in url_lower or 'img' in url_lower)
        is_bing_image = ('bing.net' in url_lower and 'th/id' in url_lower) or ('tse' in url_lower and 'bing.net' in url_lower)
        
        is_valid = has_extension or has_image_keyword or is_baidu_image or is_bing_image
        
        if not is_valid:
            print(f"❌ [URL验证] URL未通过验证: {url[:50]}...")
            print(f"   - 有扩展名: {has_extension}")
            print(f"   - 有图片关键词: {has_image_keyword}")
            print(f"   - 是百度图片: {is_baidu_image}")
            print(f"   - 是Bing图片: {is_bing_image}")
        else:
            print(f"✅ [URL验证] URL通过验证: {url[:50]}...")
            
        return is_valid

class ImageSearchWidget(QWidget):
    """图片搜索组件"""
    
    # 信号定义
    search_started = pyqtSignal(str, str)  # 查询, 搜索引擎
    search_finished = pyqtSignal(list, str)  # 结果, 错误信息
    search_progress = pyqtSignal(int)  # 进度
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.search_worker = None
        self.setup_ui()
        self.load_config()
        
    def setup_ui(self):
        """设置UI"""
        layout = QVBoxLayout(self)
        
        # 搜索设置组
        self.setup_search_settings(layout)
        
        # 搜索控制组
        self.setup_search_controls(layout)
        
        # 分页控制组
        self.setup_pagination_controls(layout)
        
    def setup_search_settings(self, parent_layout):
        """设置搜索设置"""
        settings_group = QGroupBox("搜索设置")
        settings_layout = QGridLayout(settings_group)
        
        # 搜索引擎选择
        settings_layout.addWidget(QLabel("搜索引擎:"), 0, 0)
        self.engine_combo = QComboBox()
        self.engine_combo.addItems([
            "百度图片", "Bing Images"
        ])
        settings_layout.addWidget(self.engine_combo, 0, 1)
        
        # 每页图片数量
        settings_layout.addWidget(QLabel("每页数量:"), 1, 0)
        self.per_page_spin = QSpinBox()
        self.per_page_spin.setRange(10, 100)
        self.per_page_spin.setValue(20)
        settings_layout.addWidget(self.per_page_spin, 1, 1)
        
        # 图片质量
        settings_layout.addWidget(QLabel("图片质量:"), 2, 0)
        self.quality_combo = QComboBox()
        self.quality_combo.addItems(["低", "中", "高"])
        self.quality_combo.setCurrentText("中")
        settings_layout.addWidget(self.quality_combo, 2, 1)
        
        # 图片尺寸
        settings_layout.addWidget(QLabel("最小尺寸:"), 3, 0)
        self.size_combo = QComboBox()
        self.size_combo.addItems(["任意", "小 (640x480)", "中 (1024x768)", "大 (1600x1200)", "超大 (2048x1536)"])
        settings_layout.addWidget(self.size_combo, 3, 1)
        
        # 颜色过滤
        settings_layout.addWidget(QLabel("颜色:"), 4, 0)
        self.color_combo = QComboBox()
        self.color_combo.addItems(["任意", "彩色", "黑白", "透明"])
        settings_layout.addWidget(self.color_combo, 4, 1)
        
        # 类型过滤
        settings_layout.addWidget(QLabel("类型:"), 5, 0)
        self.type_combo = QComboBox()
        self.type_combo.addItems(["任意", "照片", "插图", "GIF", "剪贴画"])
        settings_layout.addWidget(self.type_combo, 5, 1)
        
        parent_layout.addWidget(settings_group)
        
    def setup_search_controls(self, parent_layout):
        """设置搜索控制"""
        controls_group = QGroupBox("搜索控制")
        controls_layout = QVBoxLayout(controls_group)
        
        # 搜索框
        search_layout = QHBoxLayout()
        self.search_input = QLineEdit()
        self.search_input.setPlaceholderText("输入搜索关键词...")
        self.search_input.returnPressed.connect(self.start_search)
        search_layout.addWidget(self.search_input)
        
        self.search_button = QPushButton("搜索")
        self.search_button.clicked.connect(self.start_search)
        search_layout.addWidget(self.search_button)
        
        controls_layout.addLayout(search_layout)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        controls_layout.addWidget(self.progress_bar)
        
        # 状态标签
        self.status_label = QLabel("就绪")
        controls_layout.addWidget(self.status_label)
        
        parent_layout.addWidget(controls_group)
        
    def setup_pagination_controls(self, parent_layout):
        """设置分页控制"""
        pagination_group = QGroupBox("分页控制")
        pagination_layout = QHBoxLayout(pagination_group)
        
        # 上一页按钮
        self.prev_button = QPushButton("◀ 上一页")
        self.prev_button.clicked.connect(self.prev_page)
        self.prev_button.setEnabled(False)
        pagination_layout.addWidget(self.prev_button)
        
        # 页码显示
        self.page_label = QLabel("第 1 页")
        pagination_layout.addWidget(self.page_label)
        
        # 下一页按钮
        self.next_button = QPushButton("下一页 ▶")
        self.next_button.clicked.connect(self.next_page)
        self.next_button.setEnabled(False)
        pagination_layout.addWidget(self.next_button)
        
        # 跳转控制
        pagination_layout.addWidget(QLabel("跳转到:"))
        self.page_input = QSpinBox()
        self.page_input.setRange(1, 1000)
        self.page_input.setValue(1)
        pagination_layout.addWidget(self.page_input)
        
        self.go_button = QPushButton("跳转")
        self.go_button.clicked.connect(self.go_to_page)
        pagination_layout.addWidget(self.go_button)
        
        # 总页数显示
        self.total_pages_label = QLabel("共 0 页")
        pagination_layout.addWidget(self.total_pages_label)
        
        pagination_layout.addStretch()
        
        parent_layout.addWidget(pagination_group)
        
    def start_search(self):
        """开始搜索"""
        query = self.search_input.text().strip()
        if not query:
            return
            
        # 重置分页状态
        self.current_page = 1
        self.total_pages = 0
        
        # 更新UI状态
        self.search_button.setEnabled(False)
        self.progress_bar.setVisible(True)
        self.progress_bar.setValue(0)
        self.status_label.setText(f"正在搜索: {query}")
        
        # 启动搜索线程
        self.search_worker = ImageSearchWorker(
            query=query,
            search_engine=self.engine_combo.currentText(),
            page=self.current_page,
            per_page=self.per_page_spin.value()
        )
        
        self.search_worker.search_finished.connect(self.on_search_finished)
        self.search_worker.progress_updated.connect(self.on_search_progress)
        self.search_worker.start()
        
        self.search_started.emit(query, self.engine_combo.currentText())
        
    def on_search_finished(self, images, error):
        """搜索完成处理"""
        self.search_button.setEnabled(True)
        self.progress_bar.setVisible(False)
        
        if error:
            self.status_label.setText(f"搜索失败: {error}")
            self.search_finished.emit([], error)
            return
            
        # 更新分页状态
        self.total_pages = max(1, len(images) // self.per_page_spin.value())
        self.update_pagination_ui()
        
        # 更新状态
        self.status_label.setText(f"搜索完成，找到 {len(images)} 张图片")
        self.search_finished.emit(images, "")
        
    def on_search_progress(self, progress):
        """搜索进度更新"""
        self.progress_bar.setValue(progress)
        self.search_progress.emit(progress)
        
    def prev_page(self):
        """上一页"""
        if self.current_page > 1:
            self.current_page -= 1
            self.perform_search()
            
    def next_page(self):
        """下一页"""
        if self.current_page < self.total_pages:
            self.current_page += 1
            self.perform_search()
            
    def go_to_page(self):
        """跳转到指定页"""
        page = self.page_input.value()
        if 1 <= page <= self.total_pages:
            self.current_page = page
            self.perform_search()
            
    def perform_search(self):
        """执行搜索"""
        query = self.search_input.text().strip()
        if not query:
            return
            
        # 更新UI状态
        self.search_button.setEnabled(False)
        self.progress_bar.setVisible(True)
        self.progress_bar.setValue(0)
        self.status_label.setText(f"正在搜索第 {self.current_page} 页...")
        
        # 启动搜索线程
        self.search_worker = ImageSearchWorker(
            query=query,
            search_engine=self.engine_combo.currentText(),
            page=self.current_page,
            per_page=self.per_page_spin.value()
        )
        
        self.search_worker.search_finished.connect(self.on_search_finished)
        self.search_worker.progress_updated.connect(self.on_search_progress)
        self.search_worker.start()
        
    def update_pagination_ui(self):
        """更新分页UI"""
        self.page_label.setText(f"第 {self.current_page} 页")
        self.total_pages_label.setText(f"共 {self.total_pages} 页")
        self.page_input.setMaximum(self.total_pages)
        self.page_input.setValue(self.current_page)
        
        self.prev_button.setEnabled(self.current_page > 1)
        self.next_button.setEnabled(self.current_page < self.total_pages)
        
    def get_search_params(self):
        """获取搜索参数"""
        return {
            'engine': self.engine_combo.currentText(),
            'per_page': self.per_page_spin.value(),
            'quality': self.quality_combo.currentText(),
            'size': self.size_combo.currentText(),
            'color': self.color_combo.currentText(),
            'type': self.type_combo.currentText()
        }
        
    def set_search_params(self, params):
        """设置搜索参数"""
        self.engine_combo.setCurrentText(params.get('engine', '百度图片'))
        self.per_page_spin.setValue(params.get('per_page', 20))
        self.quality_combo.setCurrentText(params.get('quality', '中'))
        self.size_combo.setCurrentText(params.get('size', '任意'))
        self.color_combo.setCurrentText(params.get('color', '任意'))
        self.type_combo.setCurrentText(params.get('type', '任意'))
        
    def save_config(self):
        """保存配置"""
        config = {
            'search_params': self.get_search_params(),
            'last_query': self.search_input.text()
        }
        
        config_path = os.path.expanduser("~/.yolo_dataset_tool/image_search_config.json")
        os.makedirs(os.path.dirname(config_path), exist_ok=True)
        
        with open(config_path, 'w', encoding='utf-8') as f:
            json.dump(config, f, indent=2, ensure_ascii=False)
            
    def load_config(self):
        """加载配置"""
        config_path = os.path.expanduser("~/.yolo_dataset_tool/image_search_config.json")
        
        if os.path.exists(config_path):
            try:
                with open(config_path, 'r', encoding='utf-8') as f:
                    config = json.load(f)
                    
                self.set_search_params(config.get('search_params', {}))
                self.search_input.setText(config.get('last_query', ''))
                
            except Exception as e:
                print(f"加载图片搜索配置失败: {e}")
